package org.elasticsearch.aggregate;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.util.Counter;
import org.elasticsearch.search.aggregations.*;
import org.elasticsearch.search.aggregations.metrics.MetricsAggregator;
import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
import org.elasticsearch.search.aggregations.support.ValuesSource;
import org.elasticsearch.search.aggregations.support.ValuesSourceAggregatorFactory;
import org.elasticsearch.search.aggregations.support.ValuesSourceConfig;
import org.elasticsearch.search.internal.SearchContext;

import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicLong;

public class TimeLimitingAggregator extends MetricsAggregator {

    private final static Logger logger = LogManager.getLogger(TimeLimitingAggregator.class);

    private Counter counter;

    protected TimeLimitingAggregator(String name, SearchContext context, Aggregator parent, List<PipelineAggregator> pipelineAggregators, Map<String, Object> metaData) throws IOException {
        super(name, context, parent, pipelineAggregators, metaData);
        this.counter = context.timeEstimateCounter();
    }

    @Override
    protected LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub) throws IOException {
        final long startTime = counter.get();
        AtomicLong total = new AtomicLong(0);
        logger.warn("new");
        return new LeafBucketCollectorBase(sub, null) {
            @Override
            public void collect(int doc, long bucket) throws IOException {
                long time = counter.get();
                logger.warn("add");
                total.addAndGet(1);
                if (time - startTime - 200 > 0L || total.get() > 5) {
                    throw new RuntimeException("Time over: " + (time - startTime));
                }
            }
        };
    }

    @Override
    public InternalAggregation buildAggregation(long l) throws IOException {
        logger.warn("new internal"); // 每个分片执行一次
        return new InternalTimeLimiting(name, pipelineAggregators(), metaData());
    }

    @Override
    public InternalAggregation buildEmptyAggregation() {
        return new InternalTimeLimiting(name, pipelineAggregators(), metaData());
    }
    public static class Factory extends ValuesSourceAggregatorFactory<ValuesSource.Numeric, Factory> {

        protected Factory(
                String name, ValuesSourceConfig<ValuesSource.Numeric> config,
                SearchContext context, AggregatorFactory<?> parent,
                AggregatorFactories.Builder subFactoriesBuilder,
                Map<String, Object> metaData) throws IOException {
            super(name, config, context, parent, subFactoriesBuilder, metaData);
        }

        @Override
        protected Aggregator createUnmapped(
                Aggregator parent, List<PipelineAggregator> list,
                Map<String, Object> metaData) throws IOException {
            return new TimeLimitingAggregator(name, context, parent, list, metaData);
        }

        @Override
        protected Aggregator doCreateInternal(ValuesSource.Numeric bytes, Aggregator aggregator, boolean b, List<PipelineAggregator> list, Map<String, Object> map) throws IOException {
            return new TimeLimitingAggregator(name, context, aggregator, list, metaData);
        }
    }
}
